
import os
import sys
import json
import hmac
import time
import uuid
import base64
import hashlib
import textwrap
import tiktoken
from datetime import datetime
from email.utils import formatdate
from pydantic import BaseModel, Field
from urllib.parse import urlencode, urlparse
from typing import Dict, List, Literal, Optional, Union
from websockets.sync.client import connect as ws_connect


#  https://www.xfyun.cn/doc/spark/Web.html#_1-%E6%8E%A5%E5%8F%A3%E8%AF%B4%E6%98%8E
#  https://github.com/suqingdong/sparkapi/tree/main

MODEL_MAP = {
    "v1.5": {
        "domain": "general",
        "url": "wss://spark-api.xf-yun.com/v1.1/chat",
    },
    "v2.0": {
        "domain": "generalv2",
        "url": "wss://spark-api.xf-yun.com/v2.1/chat",
    },
    "v3.0": {
        "domain": "generalv3",
        "url": "wss://spark-api.xf-yun.com/v3.1/chat",
    },
    "v3.5": {
        "domain": "generalv3.5",
        "url": "wss://spark-api.xf-yun.com/v3.5/chat",
    },
}


def generate_rfc1123_date():
    """
    Generate a RFC 1123 compliant date string.

    """
    current_datetime = datetime.now()
    timestamp = current_datetime.timestamp()
    return formatdate(timeval=timestamp, localtime=False, usegmt=True)

def get_wss_url(api_url, api_secret, api_key):
    """
    Generate auth params for API request.
    """
    api_host = urlparse(api_url).netloc
    api_path = urlparse(api_url).path

    # step1: generate signature
    rfc1123_date = generate_rfc1123_date()
    signature_origin = textwrap.dedent(
        f"""
        host: {api_host}
        date: {rfc1123_date}
        GET {api_path} HTTP/1.1
    """
    ).strip()
    signature_sha = hmac.new(
        api_secret.encode(), signature_origin.encode(), digestmod=hashlib.sha256
    ).digest()
    signature_sha_base64 = base64.b64encode(signature_sha).decode()

    # step2: generate authorization
    authorization_payload = {
        "api_key": api_key,
        "algorithm": "hmac-sha256",
        "headers": "host date request-line",
        "signature": signature_sha_base64,
    }
    authorization_origin = ", ".join(
        f'{k}="{v}"' for k, v in authorization_payload.items()
    )
    authorization = base64.b64encode(authorization_origin.encode()).decode()

    # step3: generate wss url
    payload = {"authorization": authorization, "date": rfc1123_date, "host": api_host}
    url = api_url + "?" + urlencode(payload)
    return url

class SparkAPI(object):
    def __init__(
        self, app_id: str, api_key: str, api_secret: str, api_model: str, **kwargs
    ):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.api_model = api_model

    def create_wss_connection(self):
        api_url = MODEL_MAP[self.api_model]["url"]
        wss_url = get_wss_url(api_url, self.api_secret, self.api_key)
        return ws_connect(wss_url)

    def build_query(self, messages, **kwargs):
        query = {
            "header": {"app_id": self.app_id, "uid": str(kwargs.get("uid", 110))},
            "parameter": {
                "chat": {
                    "domain": MODEL_MAP[self.api_model]["domain"],
                    "temperature": kwargs.get("temperature", 0.5),
                    "max_tokens": kwargs.get("max_tokens", 1024),
                    "top_k": kwargs.get("top_k", 4),
                    "chat_id": str(kwargs.get("chat_id", uuid.uuid1())),
                }
            },
            "payload": {"message": {"text": messages}},
        }

        return json.dumps(query, ensure_ascii=False)

    def get_completion(self, prompt: str, **kwargs):
        """get completion from prompt"""
        messages = [{"role": "user", "content": prompt}]
        return self.get_completion_from_messages(messages, **kwargs)

    def get_resp_from_messages(self, messages: List[dict], **kwargs):
        wss = self.create_wss_connection()

        query = self.build_query(messages, **kwargs)
        wss.send(query)
        cnt = 1
        while True:
            res = json.loads(wss.recv())
            yield res
            cnt += 1
            if res["header"]["status"] == 2:
                break

    def get_completion_from_messages(self, messages: List[dict], **kwargs):

        result = self.get_resp_from_messages(messages, **kwargs)
        for res in result:
            yield res

class ChatMessage(BaseModel):
    role: Literal["user", "assistant", "system", "function"]
    content: Optional[str]
    function_call: Optional[Dict] = None

class DeltaMessage(BaseModel):
    role: Optional[Literal["user", "assistant", "system"]] = None
    content: Optional[str] = None

class ChatCompletionResponseChoice(BaseModel):
    index: int
    message: ChatMessage
    # finish_reason: Literal["stop", "length", "function_call"]
    finish_reason: Optional[str]

class ChatCompletionResponseStreamChoice(BaseModel):
    index: int
    delta: DeltaMessage
    # finish_reason: Optional[Literal["stop", "length"]]
    finish_reason: Optional[str]

class Usage(BaseModel):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

class ChatCompletionResponse(BaseModel):
    id: str = f"chatcmpl-{str(time.time())}"
    model: str
    object: Literal["chat.completion", "chat.completion.chunk"]
    choices: list
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
    usage: Optional[dict] = None

def num_tokens_from_string(string: str, encoding_name: str = "cl100k_base") -> int:
    """Returns the number of tokens in a text string."""
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens

def completion_to_openai_stream_response(completion: str, model: str = "default", index = 0, **kargs):
    completion_tokens = kargs.get("completion_tokens")
    if completion_tokens is None:
        completion_tokens = num_tokens_from_string(completion)
    prompt_tokens = kargs.get("prompt_tokens", 0)
    total_tokens = prompt_tokens + completion_tokens
    id = kargs.get("id", f"chatcmpl-{str(time.time())}")
    finish_reason = kargs.get("finish_reason", "stop")
    created = kargs.get("created", int(time.time()))
    index = kargs.get("index", 0)
    if finish_reason is None:
        usage = None
    else:
        usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
        }

    openai_response = {
        "id": id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "usage": usage,
        "choices": [
            {
                "delta": {
                    "role": "assistant",
                    "content": completion,
                },
                "index": index,
                "finish_reason": finish_reason,
            }
        ],
    }
    return openai_response

def client_response_2_chatgpt_response_stream(resp_json):
    completion = resp_json["payload"]["choices"]["text"][0]["content"]
    prompt_tokens = 0
    completion_tokens = 0
    total_tokens = 0
    if resp_json["payload"]["choices"]["status"] == 2:
        usage = resp_json["payload"]["usage"]["text"]
        prompt_tokens = usage["prompt_tokens"]
        completion_tokens = usage["completion_tokens"]
        total_tokens = usage["total_tokens"]
    id = resp_json["header"]["sid"]
    finish_reason = (
        "stop" if resp_json["payload"]["choices"]["status"] == 2 else None
    )
    return completion_to_openai_stream_response(
        completion,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        id=id,
        finish_reason=finish_reason,
    )

def to_stream(iter_content):
    for line in iter_content:
        code = line["header"]["code"]
        if code != 0:
            raise Exception(f"请求失败:{line}")
        openai_response = client_response_2_chatgpt_response_stream(line)
        yield ChatCompletionResponse(**openai_response)


if __name__ == "__main__":
    app_id = 'ac66cbc5'
    api_key = '764f6c7c69e24b95704b5fe61e14503a'
    api_secret = 'Y2Q5Zjg3ZWU3YjZlNjE1M2M4Yjc0OTJh'
    print(app_id, api_key, api_secret)
    api = SparkAPI(app_id, api_key, api_secret, "v3.5")
    messages = [{"role": "user", "content": "你好"}]
    res = api.get_resp_from_messages(messages)

    for r in to_stream(res):
        print(dict(r))